import json
import logging
from channels.generic.websocket import AsyncWebsocketConsumer
from channels.db import database_sync_to_async
from django.contrib.auth.models import AnonymousUser
from rest_framework_simplejwt.tokens import UntypedToken
from rest_framework_simplejwt.exceptions import InvalidToken, TokenError
from django.contrib.auth import get_user_model
from django.utils import timezone
from .models import Host

logger = logging.getLogger(__name__)
User = get_user_model()


class HostStatusConsumer(AsyncWebsocketConsumer):
    """主机状态更新WebSocket消费者"""
    
    async def connect(self):
        """WebSocket连接"""
        # 验证JWT token
        token = self.scope['query_string'].decode().split('token=')[-1] if 'token=' in self.scope['query_string'].decode() else None
        
        if not token:
            logger.warning("Host status WebSocket connection rejected: No token provided")
            await self.close()
            return
        
        try:
            # 验证token
            UntypedToken(token)
            # 这里可以进一步验证用户权限
            
            # 加入主机状态更新组
            await self.channel_layer.group_add(
                "host_status_updates",
                self.channel_name
            )
            
            await self.accept()
            logger.info("Host status WebSocket connection accepted")
            
        except (InvalidToken, TokenError) as e:
            logger.warning(f"Host status WebSocket connection rejected: Invalid token - {e}")
            await self.close()
    
    async def disconnect(self, close_code):
        """WebSocket断开连接"""
        # 离开主机状态更新组
        await self.channel_layer.group_discard(
            "host_status_updates",
            self.channel_name
        )
        logger.info(f"Host status WebSocket disconnected: {close_code}")
    
    async def receive(self, text_data):
        """接收WebSocket消息"""
        try:
            data = json.loads(text_data)
            message_type = data.get('type')
            
            if message_type == 'ping':
                # 心跳响应
                await self.send(text_data=json.dumps({
                    'type': 'pong',
                    'timestamp': data.get('timestamp', timezone.now().timestamp())
                }))
                
        except json.JSONDecodeError:
            logger.error(f"Invalid JSON received: {text_data}")
        except Exception as e:
            logger.error(f"Error handling message: {e}")
    
    async def host_status_update(self, event):
        """发送主机状态更新消息"""
        await self.send(text_data=json.dumps({
            'type': 'host_status_update',
            'host_id': event['host_id'],
            'hostname': event['hostname'],
            'status': event['status'],
            'last_check_time': event['last_check_time'],
            'timestamp': timezone.now().isoformat()
        }))


# 用于广播主机状态更新的辅助函数
async def broadcast_host_status_update(host_id, hostname, status, last_check_time):
    """广播主机状态更新"""
    from channels.layers import get_channel_layer
    
    channel_layer = get_channel_layer()
    if channel_layer:
        await channel_layer.group_send(
            "host_status_updates",
            {
                'type': 'host_status_update',
                'host_id': host_id,
                'hostname': hostname,
                'status': status,
                'last_check_time': last_check_time.isoformat() if last_check_time else None
            }
        )
